#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
检查PyTorch GPU支持情况
"""

def check_pytorch_gpu():
    try:
        import torch
        print("=" * 60)
        print("PyTorch和GPU支持检查")
        print("=" * 60)
        
        print(f"PyTorch版本: {torch.__version__}")
        
        # 检查CUDA可用性
        cuda_available = torch.cuda.is_available()
        print(f"CUDA可用性: {cuda_available}")
        
        if cuda_available:
            print(f"CUDA版本: {torch.version.cuda}")
            print(f"GPU数量: {torch.cuda.device_count()}")
            if torch.cuda.device_count() > 0:
                print(f"当前GPU: {torch.cuda.get_device_name()}")
                
                # 测试GPU张量操作
                print("\n测试GPU张量操作...")
                try:
                    # 创建张量并移动到GPU
                    x = torch.randn(100, 100)
                    print(f"CPU张量创建成功: {x.shape}")
                    
                    x_gpu = x.cuda()
                    print(f"GPU张量移动成功: {x_gpu.device}")
                    
                    # GPU计算
                    y_gpu = torch.matmul(x_gpu, x_gpu)
                    print(f"GPU矩阵乘法成功: {y_gpu.shape}")
                    
                    # 检查结果是否在GPU上
                    if y_gpu.is_cuda:
                        print("[INFO] GPU计算验证通过")
                    else:
                        print("[ERROR] GPU计算结果未在GPU上")
                        
                    print("\n[INFO] GPU支持检查完成!")
                    return True
                    
                except Exception as e:
                    print(f"[ERROR] GPU张量操作测试失败: {e}")
                    return False
            else:
                print("未检测到GPU设备")
                return False
        else:
            print("CUDA不可用，请检查PyTorch安装")
            return False
            
    except ImportError as e:
        print(f"无法导入PyTorch: {e}")
        print("请确保已正确安装PyTorch")
        return False
    except Exception as e:
        print(f"检查过程中出错: {e}")
        return False

if __name__ == "__main__":
    check_pytorch_gpu()